
class DataIter(object):
    def __init__(self, batches, batch_size, split=True):
        self.batches = batches
        self.batch_size = batch_size
        self.n_batches = len(self.batches) // self.batch_size
        self.residue = (len(batches) % self.n_batches != 0)  # 是否为整数
        self.index = 0
        self.split = split

    def _get_a_batch(self, b_id, e_id):  # 将数据转化成tensor
        if self.split:
            d_p = []
            batch_data = self.batches[b_id: e_id]
            for i, _ in enumerate(range(len(self.batches[0]))):
                d_p.append([item[i] for item in batch_data])
            return tuple(d_p)
        else:
            return self.batches[b_id: e_id]

    def __next__(self):
        if (self.index == self.n_batches) and (self.residue == True):
            batch_data = self._get_a_batch(self.index*self.batch_size, len(self.batches))
            self.index += 1
            return batch_data
        elif self.index >= self.n_batches:
            self.index = 0
            raise StopIteration
        else:
            batch_data = self._get_a_batch(self.index * self.batch_size,
                                           (self.index + 1) * self.batch_size)
            self.index += 1
            return batch_data

    def __iter__(self):
        return self

    def __len__(self):
        if self.residue:
            return self.n_batches + 1
        else:
            return self.n_batches



